import pandas as pd
import cv2
from PIL import Image
import os
from yolo import YOLO
import json

if __name__ == "__main__":
    yolo = YOLO()
    result = []
    dataframe = pd.read_csv("model_data/2_testa_user.csv")
    if not os.path.exists("./result/2_test_imagesa"):
        os.makedirs("./result/2_test_imagesa")
    for index in range(dataframe.shape[0]):
        # 如果遇到了错误可以下面代码直接跳过错误的那个index
        # if index<153:
        #     continue
        name = dataframe.iloc[index].values[0]
        imagepath = os.path.join("img/", name)
        try:
            image = cv2.imread(imagepath)
            image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        except:
            print('Open Error! Try again!')
            continue
        else:
            image = yolo.detect_image(image, index, result)
            # 将图片保存到result文件夹下
            image.save("./result/"+name)
    with open("result.txt", "w") as f:
        f.write(str(result))
        f.close()
    # 打包成json
    with open("result.json", "w") as file:
        json.dump(result, file)
        file.close()
    print("结束了")
